import pymysql
import pymysql.cursors
import json
from config_manager import config


def get_db_connection():
    """创建数据库连接"""
    try:
        conn = pymysql.connect(
            host=config['database']['host'],
            port=config['database']['port'],
            user=config['database']['user'],
            password=config['database']['password'],
            database=config['database']['database'],
            cursorclass=pymysql.cursors.DictCursor
        )
        return conn
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return None


def create_tables():
    """创建数据库表（如果不存在）"""
    conn = get_db_connection()
    if not conn:
        print("无法连接数据库，跳过表创建")
        return
    
    try:
        with conn.cursor() as cursor:
            # 创建股票分析表
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS stock_analysis_results (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    symbol VARCHAR(20) NOT NULL,
                    stock_name VARCHAR(100),
                    industry VARCHAR(100),
                    indicators JSON,
                    analysis TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    INDEX idx_symbol (symbol),
                    INDEX idx_created_at (created_at)
                )
            """)
            
            # 创建股票信息表，用于存储所有股票的代码、名称和行业映射
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS stock_info (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    symbol VARCHAR(20) NOT NULL,
                    symbol_standardized VARCHAR(10) NOT NULL,  # 标准化的股票代码（去掉后缀）
                    stock_name VARCHAR(100) NOT NULL,
                    industry VARCHAR(100),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    UNIQUE KEY uk_symbol (symbol),
                    INDEX idx_symbol_standardized (symbol_standardized)
                )
            """)
        conn.commit()
        print("数据库表创建/检查成功")
        
        # 初始化股票数据
        initialize_stock_data(conn)
    except Exception as e:
        print(f"创建表失败: {e}")
    finally:
        conn.close()


def initialize_stock_data(conn):
    """初始化股票数据"""
    try:
        with conn.cursor() as cursor:
            # 检查是否已有数据
            cursor.execute("SELECT COUNT(*) FROM stock_info")
            count = cursor.fetchone()['COUNT(*)']
            
            if count == 0:
                # 插入一些基本的股票数据
                initial_stocks = [
                    ('600519.SS', '600519', '贵州茅台', '白酒'),
                    ('000858.SZ', '000858', '五粮液', '白酒'),
                    ('002594.SZ', '002594', '比亚迪', '新能源汽车'),
                    ('601318.SS', '601318', '中国平安', '保险'),
                    ('600036.SS', '600036', '招商银行', '银行'),
                    ('000002.SZ', '000002', '万科A', '房地产'),
                    ('600000.SS', '600000', '浦发银行', '银行'),
                    ('000725.SZ', '000725', '京东方A', '电子'),
                    ('601166.SS', '601166', '兴业银行', '银行'),
                    ('601328.SS', '601328', '交通银行', '银行'),
                    ('600031.SS', '600031', '三一重工', '工程机械'),
                    ('601668.SS', '601668', '中国建筑', '建筑装饰'),
                    ('600019.SS', '600019', '宝钢股份', '钢铁'),
                    ('000858.SZ', '000858', '五粮液', '白酒') # 不带后缀的格式
                ]
                
                # 使用批量插入
                sql = """
                    INSERT IGNORE INTO stock_info 
                    (symbol, symbol_standardized, stock_name, industry) 
                    VALUES (%s, %s, %s, %s)
                """
                cursor.executemany(sql, initial_stocks)
                conn.commit()
                print(f"成功初始化 {cursor.rowcount} 条股票数据")
    except Exception as e:
        print(f"初始化股票数据失败: {e}")
        conn.rollback()
